# import os
#
# os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'  # 临时解决OMP冲突
#
# from ultralytics import YOLO
#
# if __name__ == '__main__':
#     model = YOLO('yolov8n.yaml')
#     results = model.train(
#         data='mydata.yaml',
#         epochs=200,
#         batch=8,           # 根据显存调整
#         workers=0,         # Windows下避免多进程问题
#         amp=False,         # 若AMP失败则禁用
#         device='0'
#     )

# import os
# import shutil
# import warnings
# import pandas as pd
# import matplotlib.pyplot as plt
# from ultralytics import YOLO
# from pathlib import Path
#
# # ========== 1. 环境设置 ==========
# warnings.filterwarnings('ignore')  # 禁用警告提示
# os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"  # 解决Windows环境冲突[^3]
#
# # ========== 2. 创建结果保存目录 ==========
# results_dir = Path("./training_results")
# results_dir.mkdir(exist_ok=True, parents=True)
# print(f"✅ 结果保存目录: {results_dir.resolve()}")
#
# # ========== 3. 加载预训练模型 ==========
# model = YOLO('yolov8n.pt')  # 官方预训练模型
# # model = YOLO('custom_pretrained.pt')  # 自定义预训练权重
#
# # ========== 4. 训练配置 ==========
# train_params = {
#     'data': 'mydata.yaml',  # 数据集配置文件[^2]
#     'epochs': 100,  # 训练轮数
#     'batch': 16,  # 批次大小
#     'imgsz': 640,  # 输入尺寸
#     'optimizer': 'SGD',  # 推荐优化器(SGD优于AdamW)[^1]
#     'lr0': 0.001,  # 初始学习率($\eta_0$)
#     'momentum': 0.9,  # 动量参数($\gamma$)
#     'weight_decay': 0.0005,  # 权重衰减
#     'device': '0',  # GPU设备ID
#     'project': str(results_dir),  # 结果保存路径
#     'name': 'exp_smoking_det',  # 实验名称
#     'exist_ok': True,  # 覆盖已有结果
#     'save_period': 10,  # 每10轮保存一次权重
#     'patience': 30,  # 早停机制阈值
#     'close_mosaic': 15,  # 关闭Mosaic增强的轮次
#     'workers': 8,  # 数据加载线程数
#     'plots': True,  # 自动生成训练曲线
# }
#
# # ========== 5. 执行训练 ==========
# results = model.train(**train_params)
#
#
# # ========== 6. 手动保存关键指标 ==========
# def save_training_artifacts(results_dir):
#     """保存训练曲线图和CSV指标数据"""
#     # 定位结果文件
#     results_csv = next(Path(results_dir).rglob('results.csv'))
#     curves_dir = Path(results_dir) / "exp_smoking_det"
#
#     # 关键曲线文件
#     curve_files = [
#         'results.png',  # 综合指标
#         'confusion_matrix.png',  # 混淆矩阵
#         'val_batch0_pred.jpg'  # 验证样例
#     ]
#
#     # 复制曲线图到目标目录
#     for fname in curve_files:
#         src = curves_dir / fname
#         if src.exists():
#             shutil.copy(src, results_dir / fname)
#             print(f"📊 已保存曲线图: {fname}")
#
#     # 处理CSV指标数据
#     df = pd.read_csv(results_csv)
#     df = df.rename(columns={
#         'epoch': '轮次',
#         'metrics/precision(B)': '精确率',
#         'metrics/recall(B)': '召回率',
#         'metrics/mAP50(B)': 'mAP@50',
#         'metrics/mAP50-95(B)': 'mAP@50-95',
#         'train/box_loss': '边界框损失',
#         'train/cls_loss': '分类损失'
#     })
#     df.to_csv(results_dir / "training_metrics.csv", index=False)
#     print(f"📈 已保存指标CSV: training_metrics.csv")
#
#
# save_training_artifacts(results_dir)
# print("🚀 训练完成! 所有结果保存在:", results_dir.resolve())

# import os
# import shutil
# import torch
# import warnings
# from datetime import datetime
# from ultralytics import YOLO
#
# # ======== 1. 环境修复 ========
# warnings.filterwarnings('ignore')
# os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
#
#
# # ======== 2. 多进程问题修复 ========
# def fix_multiprocessing():
#     """修复Windows多进程启动问题[^3]"""
#     if os.name == 'nt':  # Windows系统
#         torch.multiprocessing.freeze_support()
#         torch.multiprocessing.set_start_method('spawn', force=True)
#         print("🛠️ 已设置多进程启动方法为'spawn'")
#
#
# if __name__ == '__main__':
#     fix_multiprocessing()  # 必须放在__main__内
#
#     # ======== 3. 模型配置 ========
#     model = YOLO('yolov8n.yaml')  # 或 'yolov8m.pt'
#
#     # ======== 4. 训练参数优化 ========
#     train_params = {
#         'data': 'mydata.yaml',
#         'epochs': 10,
#         'batch': 8,
#         'imgsz': 640,
#         'workers': 0,  # 关键修复：禁用多进程加载[^1][^4]
#         'optimizer': 'SGD',  # 降低震荡风险
#         'lr0': 0.01,  # 初始学习率($\eta_0$)
#         'momentum': 0.937,  # 动量参数($\gamma$)
#         'weight_decay': 0.0005,  # 权重衰减($\lambda$)
#         'amp': False,  # 关闭自动混合精度
#         'patience': 50,  # 早停机制阈值
#     }
#
#     # ======== 5. 结果目录配置 ========
#     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
#     results_dir = f"runs/detect/training_results_{timestamp}"
#     train_params.update({
#         'project': results_dir,
#         'name': 'train',
#         'exist_ok': True,
#         'save_period': 5
#     })
#
#     # ======== 6. 执行训练 ========
#     results = model.train(**train_params)
#     print("✅ 训练完成! 结果保存在:", os.path.abspath(results_dir))
#
#     # ======== 7. 结果文件处理 ========
#     train_results_path = os.path.join(results_dir, 'train_jujube_detection')
#     if os.path.exists(train_results_path):
#         for file in ['results.png', 'confusion_matrix.png', 'results.csv']:
#             src = os.path.join(train_results_path, file)
#             if os.path.exists(src):
#                 shutil.copy(src, results_dir)
#                 print(f"📊 已复制: {file}")

import os
import shutil
import torch
import warnings
from datetime import datetime
from ultralytics import YOLO

# ======== 1. 环境修复 ========
warnings.filterwarnings('ignore')
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'


# ======== 2. 多进程问题修复 ========
def fix_multiprocessing():
    """修复Windows多进程启动问题"""
    if os.name == 'nt':  # Windows系统
        torch.multiprocessing.freeze_support()
        torch.multiprocessing.set_start_method('spawn', force=True)
        print("🛠️ 已设置多进程启动方法为'spawn'")


if __name__ == '__main__':
    fix_multiprocessing()  # 必须放在__main__内

    # ======== 3. 模型配置（关键修改：使用分割模型） ========
    # 可选：从配置文件从头训练（yolov8n-seg.yaml 对应nano版分割模型）
    # 或从预训练权重微调（yolov8n-seg.pt，推荐，收敛更快）
    model = YOLO('yolov8n-seg.yaml')  # 分割模型配置文件（nano版，轻量快速）
    # model = YOLO('yolov8m-seg.pt')  # 若需要中等精度，可用m版预训练权重

    # ======== 4. 训练参数优化（分割任务适配） ========
    train_params = {
        'data': 'mydata_seg.yaml',  # 关键修改：分割任务的数据配置文件（需包含掩码路径）
        'epochs': 50,  # 分割任务通常需要更多迭代（建议30-100）
        'batch': 8,
        'imgsz': 640,  # 输入图片尺寸（保持与检测一致即可）
        'workers': 0,  # Windows系统禁用多进程，避免报错
        'optimizer': 'Adam',  # 分割任务对优化器更敏感，Adam收敛更稳（替代SGD）
        'lr0': 0.001,  # 初始学习率（分割任务建议更小，避免震荡）
        'momentum': 0.9,
        'weight_decay': 0.0005,
        'amp': True,  # 分割任务计算量更大，开启自动混合精度加速训练（需GPU支持）
        'patience': 30,  # 早停机制：30轮无提升则停止
        'mask_overlap': 0.5,  # 分割掩码的重叠阈值（默认即可）
    }

    # ======== 5. 结果目录配置 ========
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    results_dir = f"runs/segment/training_results_{timestamp}"  # 关键修改：结果目录放在segment下（区分检测）
    train_params.update({
        'project': results_dir,
        'name': 'train',
        'exist_ok': True,
        'save_period': 10,  # 每10轮保存一次模型
    })

    # ======== 6. 执行训练 ========
    results = model.train(** train_params)
    print("✅ 训练完成! 结果保存在:", os.path.abspath(results_dir))

    # ======== 7. 结果文件处理（复制关键指标） ========
    train_results_path = os.path.join(results_dir, 'train')
    if os.path.exists(train_results_path):
        # 分割任务关键结果文件（比检测多了掩码相关指标）
        for file in ['results.png', 'confusion_matrix.png', 'results.csv', 'val_mask_iou.png']:
            src = os.path.join(train_results_path, file)
            if os.path.exists(src):
                shutil.copy(src, results_dir)
                print(f"📊 已复制: {file}")
